import mindspore as ms
import mindspore.nn as nn


def get_accuracy(network, datasets):
    metric = nn.Accuracy('classification')
    for x, y in datasets:
        y_hat = network(x)
        metric.update(y_hat, y)
    return metric.eval()


def get_loss(network, loss_f, dataset):
    loss = 0
    for x, y in dataset:
        y_hat = network(x)
        loss += loss_f(y_hat, y)
    return loss


class AccuracyMonitor(ms.Callback):
    def __init__(self, valid_data):
        self.valid_data = valid_data

    def on_train_epoch_end(self, run_context):
        """Called after each epoch finished."""

        callback_params = run_context.original_args()

        cur_epoch_num = callback_params.cur_epoch_num
        epoch_num = callback_params.epoch_num

        network = callback_params.network

        train_data = callback_params.train_dataset

        train_accu = get_accuracy(network, train_data)
        valid_accu = get_accuracy(network, self.valid_data)

        loss = get_loss(network, callback_params.loss_fn, train_data)
        loss /= callback_params.batch_num

        print(f'epoch:[{cur_epoch_num}/{epoch_num}] Loss:{loss} Train Accuracy:{train_accu} Valid Accuracy:{valid_accu}')
